#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import ctypes
import platform

if platform.system() != "Linux":
    raise Exception("Only support linux")

from ctypes.util import find_library
_libcname = find_library("c")
libc = ctypes.CDLL(_libcname)

'''include/uapi/asm-generic/ioctl.h start'''
_IOC_NRBITS     = 8
_IOC_TYPEBITS   = 8
_IOC_SIZEBITS   = 14
_IOC_DIRBITS    = 2

_IOC_NRMASK     = ((1 << _IOC_NRBITS)-1) & 0xFFFFFFFF
_IOC_TYPEMASK   = ((1 << _IOC_TYPEBITS)-1) & 0xFFFFFFFF
_IOC_SIZEMASK   = ((1 << _IOC_SIZEBITS)-1) & 0xFFFFFFFF
_IOC_DIRMASK    = ((1 << _IOC_DIRBITS)-1) & 0xFFFFFFFF

_IOC_NRSHIFT    = 0
_IOC_TYPESHIFT  = (_IOC_NRSHIFT+_IOC_NRBITS)
_IOC_SIZESHIFT  = (_IOC_TYPESHIFT+_IOC_TYPEBITS)
_IOC_DIRSHIFT   = (_IOC_SIZESHIFT+_IOC_SIZEBITS)

_IOC_NONE       = 0
_IOC_WRITE      = 1
_IOC_READ       = 2

def _IOC(d, t, nr, size) -> int:
    if isinstance(t, str):
        assert len(t) == 1
        t = ord(t)
    nr = (d << _IOC_DIRSHIFT) |         \
            (t << _IOC_TYPESHIFT) |     \
            (nr << _IOC_NRSHIFT) |      \
            (size << _IOC_SIZESHIFT)
    return nr & 0xFFFFFFFF

_IOC_TYPECHECK  = lambda t: ctypes.sizeof(t)

_IO             = lambda t, nr: _IOC(_IOC_NONE, t, nr, 0)
_IOR            = lambda t, nr, size: _IOC(_IOC_READ, t, nr, _IOC_TYPECHECK(size))
_IOW            = lambda t, nr, size: _IOC(_IOC_WRITE, t, nr, _IOC_TYPECHECK(size))
_IOWR           = lambda t, nr, size: _IOC(_IOC_READ | _IOC_WRITE, t, nr, _IOC_TYPECHECK(size))
_IOR_BAD        = lambda t, nr, size: _IOC(_IOC_READ, t, nr, ctypes.sizeof(size))
_IOW_BAD        = lambda t, nr, size: _IOC(_IOC_WRITE, t, nr, ctypes.sizeof(size))
_IOWR_BAD       = lambda t, nr, size: _IOC(_IOC_READ | _IOC_WRITE, t, nr, ctypes.sizeof(size))


_IOC_DIR        = lambda nr: (nr >> _IOC_DIRSHIFT) & _IOC_DIRMASK
_IOC_TYPE       = lambda nr: (nr >> _IOC_TYPESHIFT) & _IOC_TYPEMASK
_IOC_NR         = lambda nr: (nr >> _IOC_NRSHIFT) & _IOC_NRMASK
_IOC_SIZE       = lambda nr: (nr >> _IOC_SIZESHIFT) & _IOC_SIZEMASK


IOC_IN          = _IOC_WRITE << _IOC_DIRSHIFT
IOC_OUT         = _IOC_READ << _IOC_DIRSHIFT
IOC_INOUT       = (_IOC_WRITE|_IOC_READ) << _IOC_DIRSHIFT
IOCSIZE_MASK    = _IOC_SIZEMASK << _IOC_SIZESHIFT
IOCSIZE_SHIFT   = _IOC_SIZESHIFT
'''include/uapi/asm-generic/ioctl.h end'''

'''export start'''
IOC_NONE        = _IOC_NONE
IOC_READ        = _IOC_READ
IOC_WRITE       = _IOC_WRITE

IOC             = _IOC
IO              = _IO
IOR             = _IOR
IOW             = _IOW
IOWR            = _IOWR
'''export end'''

class FileIoctl(object):

    def __init__(self, filepath, flags=os.O_RDWR):
        self._filepath = filepath
        try:
            self._fd = os.open(self._filepath, flags=flags)
        except Exception as e:
            raise e

    def __del__(self):
        try:
            os.close(self._fd)
        except Exception as e:
            pass

    def __enter__(self):
        return self

    def __exit__(*args):
        self.__del__()

    @property
    def filepath(self):
        return self._filepath

    def do_ioctl(self, request, data=None):
        # do not use fcntl.ioctl
        if data is None:
            return libc.ioctl(self._fd, request)

        arg = None

        if isinstance(data, (ctypes.Structure, ctypes.Union)):
            arg = ctypes.pointer(data)

        if isinstance(data, (int,)):
            arg = data

        simple_type_list = (
            ctypes.c_short, ctypes.c_ushort,
            ctypes.c_long, ctypes.c_ulong,
            ctypes.c_int, ctypes.c_uint,
            ctypes.c_longlong, ctypes.c_ulonglong,
            ctypes.c_byte, ctypes.c_ubyte,
            ctypes.c_char, ctypes.c_char,
        )
        if isinstance(data, simple_type_list):
            arg = ctypes.byref(data)

        if isinstance(data, bytes):
            # struct.pack?
            arg = data

        if arg:
            return libc.ioctl(self._fd, request, arg)

        # if isinstance(data, bytearray):
        #     return libc.ioctl(self._fd, request, data)

        raise TypeError(f"unknown type: {type(data)}")


def demo_random():
    '''
    #include <stdio.h>
    #include <unistd.h>
    #include <fcntl.h>
    #include <sys/ioctl.h>
    #include <linux/random.h>

    int main() {
        int fd;
        int entropy_count;

        // 打开/dev/random文件
        fd = open("/dev/random", O_RDONLY);
        if (fd == -1) {
            perror("open");
            return 1;
        }

        // 使用ioctl发送RNDGETENTCNT命令并获取返回的熵计数
        if (ioctl(fd, RNDGETENTCNT, &entropy_count) == -1) {
            perror("ioctl");
            close(fd);
            return 1;
        }

        printf("RNDGETENTCNT: 0x%08X\n", RNDGETENTCNT);
        printf("entropy count: %d\n", entropy_count);

        close(fd);
        return 0;
    }
    '''

    # Get the entropy count.
    RNDGETENTCNT            = IOR('R', 0x00, ctypes.c_int)

    # Add to (or subtract from) the entropy count.  (Superuser only.)
    RNDADDTOENTCNT          = IOWR('R', 0x01, ctypes.c_int)

    # Get the contents of the entropy pool.  (Superuser only.)
    RNDGETPOOL              = IOR('R', 0x02, ctypes.c_int * 2)

    # Write bytes into the entropy pool and add to the entropy count.
    # (Superuser only.)
    RNDADDENTROPY           = IOW("R", 0x03, ctypes.c_int * 2)

    # Clear entropy count to 0.  (Superuser only.)
    RNDZAPENTCNT            = IO('R', 0x04)

    # Clear the entropy pool and associated counters.  (Superuser only.)
    RNDCLEARPOOL            = IO("R", 0x06)

    # Reseed CRNG.  (Superuser only.)
    RNDRESEEDCRNG           = IO("R", 0x07)

    rand = FileIoctl("/dev/random")

    data = ctypes.c_int()
    result = rand.do_ioctl(RNDGETENTCNT, data)
    print(f"RNDGETENTCNT: 0x{RNDGETENTCNT:X}")
    print(f"random ioctl result: {result}")
    print(f"entropy count: {data.value}")

    data = b'\x00' * 4
    result = rand.do_ioctl(RNDGETENTCNT, data)
    print(f"RNDGETENTCNT: 0x{RNDGETENTCNT:X}")
    print(f"random ioctl result: {result}")
    print(f"entropy count: {int.from_bytes(data, byteorder='little')}")


def demo_termios():
    import termios

    class termios_type(ctypes.Structure):
        '''struct termios'''
        NCCS = 19
        _pack_ = 2
        _fields_ = [
            ("c_iflag", ctypes.c_ushort),
            ("c_oflag", ctypes.c_ushort),
            ("c_cflag", ctypes.c_ushort),
            ("c_lflag", ctypes.c_ushort),
            ("c_cc",    ctypes.c_char * NCCS),
        ]

    data = termios_type()
    sout = FileIoctl('/dev/stdout')
    result = sout.do_ioctl(termios.TIOCGWINSZ, data)
    print(f"termios.TIOCGWINSZ: 0x{termios.TIOCGWINSZ:X}")
    print(f"stdout ioctl result: {result}")
    print(f"data.c_iflag, data.c_oflag, data.c_cflag, data.c_lflag: "
          f"{data.c_iflag}, {data.c_oflag}, {data.c_cflag}, {data.c_lflag}")
    print(f"data.c_cc: {data.c_cc}")


def main():
    demo_random()
    print()
    demo_termios()

if __name__ == "__main__":
    main()
